import matplotlib.pyplot as plt
import networkx as nx
import random
from matplotlib.animation import FuncAnimation

def prims_algorithm_visualization(ax):
    # Generate a random connected undirected graph
    G = nx.Graph()

    # Create nodes
    num_nodes = random.randint(5, 5)
    G.add_nodes_from(range(num_nodes))

    # Add random edges with weights to make it a connected graph
    for u in range(num_nodes):
        for v in range(u + 1, num_nodes):
            if not G.has_edge(u, v):
                weight = random.randint(1, 10)
                G.add_edge(u, v, weight=weight)

    # Create positions for nodes
    pos = nx.spring_layout(G)

    # Draw initial graph
    def draw_graph(mst=False):
        ax.clear()
        node_color = 'lightblue' if not mst else 'lightgreen'
        nx.draw_networkx_nodes(G, pos, ax=ax, node_color=node_color, node_size=500)
        edge_color = 'lightgrey' if not mst else 'green'
        width = 1 if not mst else 2
        edges = G.edges() if not mst else mst_edges
        nx.draw_networkx_edges(G, pos, edgelist=edges, ax=ax, edge_color=edge_color, width=width)
        edge_labels = {(u, v): G[u][v]['weight'] for (u, v) in G.edges()}
        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, ax=ax, font_color='red')
        nx.draw_networkx_labels(G, pos, labels={node: str(node) for node in G.nodes()}, font_size=12, font_color='black', ax=ax)
        ax.set_aspect('equal')
        ax.axis('off')

    draw_graph()

    # Prim's algorithm setup
    mst_edges = []
    visited = set()
    start_node = random.choice(list(G.nodes()))
    visited.add(start_node)
    edges = [(G[start_node][neighbor]['weight'], start_node, neighbor) for neighbor in G.neighbors(start_node)]

    def update(num):
        if edges:
            # Get the edge with the minimum weight that does not form a cycle
            edges.sort()
            for i, (weight, u, v) in enumerate(edges):
                if v not in visited:
                    visited.add(v)
                    mst_edges.append((u, v))
                    # Add all new edges to the list
                    for neighbor in G.neighbors(v):
                        if neighbor not in visited:
                            edges.append((G[v][neighbor]['weight'], v, neighbor))
                    del edges[i]  # Remove the edge from the list
                    break

            # Redraw graph with the current MST edges
            draw_graph()
            nx.draw_networkx_edges(G, pos, edgelist=mst_edges, ax=ax, edge_color='green', width=2)
            ax.set_title(f'Prim\'s Algorithm Visualization: Step {num}')
        else:
            # Draw the final MST in tree layout
            ax.clear()
            T = nx.Graph()
            T.add_edges_from(mst_edges)
            pos_tree = nx.nx_agraph.graphviz_layout(T, prog='dot')  # Ensure pygraphviz is installed
            nx.draw(T, pos=pos_tree, with_labels=True, node_color='lightgreen', edge_color='green', node_size=500, width=2, ax=ax)
            edge_labels = {(u, v): G[u][v]['weight'] for (u, v) in mst_edges}
            nx.draw_networkx_edge_labels(T, pos_tree, edge_labels=edge_labels, ax=ax, font_color='red')
            ax.set_title('Final Minimum Spanning Tree')

    ani = FuncAnimation(ax.figure, update, frames=range(len(G.edges()) + 1), repeat=False, interval=1000)
    return ani